import baostock as bs
import pandas as pd
import numpy as np
from scipy import stats


def get_hs300_data(date):
    """
    获取沪深300成分股数据
    :return: 沪深300成分股数据
    """
    # 获取沪深300成分股
    rs = bs.query_hs300_stocks(date)
    hs300_stocks = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        hs300_stocks.append(rs.get_row_data())
    result = pd.DataFrame(hs300_stocks, columns=rs.fields)
    return result

def get_sh_50_data(date):
    # 获取上证50成分股
    rs = bs.query_sz50_stocks(date)
    sh50_stocks = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        sh50_stocks.append(rs.get_row_data())
    result = pd.DataFrame(sh50_stocks, columns=rs.fields)
    return result

def get_all_stock_data(date):
    """
    获取所有股票数据
    :return: 所有股票数据
    """
    # 获取所有股票
    rs = bs.query_all_stock(date)
    all_stocks = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        all_stocks.append(rs.get_row_data())
    result = pd.DataFrame(all_stocks, columns=rs.fields)
    return result

def get_code_profit_data(code, start_year, end_year):
    """
    获取股票的盈利能力数据
    :param code: 股票代码
    :param start_year: 开始年份
    :param end_year: 结束年份
    :return: 股票的盈利能力数据
    """
    # 查询季频估值指标盈利能力
    profit_list = []
    columns = None
    for year in range(start_year, end_year + 1):
        rs_profit = bs.query_profit_data(code=code, year=year, quarter=4)
        while (rs_profit.error_code == '0') & rs_profit.next():
            profit_list.append(rs_profit.get_row_data())
            if columns is None:
                columns = rs_profit.fields
    result_profit = pd.DataFrame(profit_list, columns=columns)
    return result_profit

def is_profit_stable(code, start_year, end_year):
    """
    判断股票的盈利能力是否稳定
    :param code: 股票代码
    :param start_year: 开始年份
    :param end_year: 结束年份
    :return: 股票的盈利能力是否稳定
    """
    # 查询季频估值指标盈利能力
    result_profit = get_code_profit_data(code, start_year, end_year)
    try:
        net_profit_data = result_profit['netProfit']
        # 字符串转换成小数
        net_profit_data = net_profit_data.astype(float)
        rs = is_stable_growth(net_profit_data)
        return rs['is_stable_growth']
    except Exception as err:
        print('获取盈利数据失败')
        return False


def is_stable_growth(series, significance_level=0.05, window=3, trend_threshold=0.5, volatility_threshold=0.38):
    """
    判断一个pandas.Series是否呈现统计学意义上的稳定增长。

    参数:
    series: pandas.Series - 待检验的时间序列数据
    significance_level: float - 统计检验的显著性水平（默认0.05）
    window: int - 用于计算移动平均和移动标准差的窗口大小（默认3）
    trend_threshold: float - 趋势稳定性的阈值（默认0.5）
    volatility_threshold: float - 波动率的阈值（默认0.2）

    返回:
    dict: 包含判断结果和详细统计信息的字典
    """

    # 检查输入数据有效性
    if not isinstance(series, pd.Series):
        raise TypeError("输入必须为pandas.Series对象")

    if series.empty or len(series) < 3:
        return {
            'is_stable_growth': False,
            'reason': '数据量不足，至少需要3个数据点',
            'details': {}
        }

    # 预处理：去除NaN值
    clean_series = series.dropna()
    if len(clean_series) < 3:
        return {
            'is_stable_growth': False,
            'reason': '有效数据量不足',
            'details': {}
        }

    # 计算基本增长指标
    growth_rates = clean_series.pct_change().dropna()  # 计算增长率[5](@ref)

    # 1. 趋势分析 - 使用线性回归检验趋势的显著性[3](@ref)
    x = np.arange(len(clean_series)).reshape(-1, 1)
    y = clean_series.values
    slope, intercept, r_value, p_value, std_err = stats.linregress(x.flatten(), y)

    # 2. 波动性分析 - 计算增长率的移动标准差[3](@ref)
    moving_std = growth_rates.rolling(window=window).std().dropna()
    avg_volatility = moving_std.mean() if not moving_std.empty else np.nan

    # 综合判断条件
    has_positive_trend = slope > 0 and p_value < significance_level
    has_low_volatility = not np.isnan(avg_volatility) and avg_volatility < volatility_threshold

    # 计算趋势稳定性得分 (0-1范围)
    trend_stability_score = 0
    if has_positive_trend:
        # 基于R²和波动性的综合评分
        trend_stability_score = min(r_value ** 2 * (1 - avg_volatility) / trend_threshold, 1.0)

    # 最终判断
    is_stable = (
            has_positive_trend and
            has_low_volatility and
            trend_stability_score >= trend_threshold
    )

    return {
        'is_stable_growth': is_stable,
        'details': {
            'trend_slope': slope,
            'trend_p_value': p_value,
            'r_squared': r_value ** 2,
            'avg_volatility': avg_volatility,
            'trend_stability_score': trend_stability_score,
            'growth_rates_mean': growth_rates.mean(),
            'growth_rates_std': growth_rates.std()
        },
        'diagnosis': (
            "稳定增长" if is_stable else
            "趋势不显著" if not has_positive_trend else
            "波动过大" if not has_low_volatility else
            "趋势稳定性不足"
        )
    }

def get_roe(code, start_year, end_year):
    """
    获取股票的ROE
    :param code: 股票代码
    :param start_year: 开始年份
    :param end_year: 结束年份
    :return: 股票的ROE
    """
    result_profit = get_code_profit_data(code, start_year, end_year)
    roe_data = result_profit['roeAvg'].astype(float)
    return roe_data.mean()

def get_pe(code, date):
    """
    获取股票的PE
    :param code: 股票代码
    :param date: 日期
    :return: 股票的PE
    """
    rs = bs.query_history_k_data_plus(code,
        "date,peTTM",
        frequency="d", adjustflag="3", start_date=date, end_date=date)
    result_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        result_list.append(rs.get_row_data())
    result = pd.DataFrame(result_list, columns=rs.fields)
    if len(result) == 0:
        print(f'未找到{code}在{date}的PE数据')
        raise ValueError(f'未找到{code}在{date}的PE数据')
    return result['peTTM'].astype(float)[0]

def get_close_price(code, date):
    """
    获取股票的收盘价
    :param code: 股票代码
    :param date: 日期
    :return: 股票的收盘价
    """
    rs = bs.query_history_k_data_plus(code, "code,close", start_date=date,
                                      end_date=date, frequency="d", adjustflag="3")
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)
    if len(result) == 0:
        print(f'未找到{code}在{date}的收盘价数据')
        raise ValueError(f'未找到{code}在{date}的收盘价数据')
    return float(result['close'][0])

def to_rate(data):
    """
    转换为收益率
    :param data: 数据
    :return: 收益率
    """
    # 转换为收益率
    profit_rate = [{'date': data[0]['date'], 'rate': 0}]
    initial_amount = data[0]['value']
    for i in range(1, len(data)):
        rate = (data[i]['value'] - initial_amount) / initial_amount
        profit_rate.append({'date': data[i]['date'], 'rate': rate})
    return profit_rate

def get_pr(code, date):
    """
    获取股票的PR
    :param code: 股票代码
    :param date: 日期
    :return: 股票的PR
    """
    pe = get_pe(code, date)
    end_year = int(date[:4])
    avg_roe = get_roe(code, end_year - 4, end_year)
    return pe / (avg_roe * 100)

def get_stable_code_from_hs300(year):
    """
    获取沪深300中稳定增长的股票
    :param year: 年份
    :return: 稳定增长的股票代码列表
    """
    print(f'获取{year}年的沪深300中表现稳定的股票')
    first_workday_of_year = str(get_first_workday_of_year(year))
    hs300_data = get_hs300_data(first_workday_of_year)
    stable_code_list = []
    i = 1
    for index, row in hs300_data.iterrows():
        # 判断股票的盈利能力是否稳定
        if is_profit_stable(row['code'], year - 4, year):
            # print(f'{i}/300. {row["code_name"]} is stable')
            stable_code_list.append(row)
        else:
            # print(f'{i}/300. {row["code_name"]} is NOT stable')
            pass
        i += 1
    return pd.DataFrame(stable_code_list)

def save_as_excel(df, file_path):
    """
    保存DataFrame到Excel文件
    :param df: 要保存的DataFrame
    :param file_path: 保存的文件路径
    """
    with pd.ExcelWriter(file_path, engine='xlsxwriter') as writer:
        df.to_excel(writer, sheet_name='Sheet1', index=False, startrow=1)
        # Get the workbook and worksheet objects
        workbook = writer.book
        worksheet = writer.sheets['Sheet1']
        # Define a format for the header
        header_format = workbook.add_format({'bold': True, 'bg_color': '#FFFF00', 'border': 1})
        # Apply the format to the header row
        for col_num, value in enumerate(df.columns.values):
            worksheet.write(0, col_num + 1, value, header_format)
